import os
from dotenv import load_dotenv
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain.memory import ConversationBufferMemory
from langchain.chains import LLMChain

# 导入我们自定义的智谱AI LLM
from llm_zhipu import ZhipuAIChatLLM

# 加载环境变量
load_dotenv()

# 初始化LLM - 使用我们自定义的智谱AI包装器
llm = ZhipuAIChatLLM(
    model_name="glm-4-flash",
    temperature=0.7,
    max_tokens=2048,
    api_key=os.getenv("ZHIPU_API_KEY")
)

# 创建记忆 - 保存对话历史
memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)

# 创建提示模板 - 包含对话历史
prompt = ChatPromptTemplate.from_messages([
    ("system", "你是一个友好的个人助手，回答要自然、有帮助。"),
    MessagesPlaceholder(variable_name="chat_history"),  # 这里会插入历史对话
    ("human", "{input}")
])

# 构建链
chain = LLMChain(
    llm=llm,
    prompt=prompt,
    memory=memory,
    verbose=True  # 显示详细执行过程
)

def chat_with_assistant():
    print("个人助手已启动! 输入'退出'或'quit'结束对话")
    print("-" * 50)
    
    while True:
        user_input = input("\n你: ").strip()
        
        if user_input.lower() in ['退出', 'quit', 'exit']:
            print("助手: 再见！很高兴为您服务！")
            break
            
        if not user_input:
            continue
            
        try:
            # 调用链并获取响应
            response = chain.invoke({"input": user_input})
            print(f"助手: {response['text']}")
            
        except Exception as e:
            print(f"出错: {e}")

if __name__ == "__main__":
    chat_with_assistant()